#include "Graphics/PiccoloVKRenderPass.h"

#include <Graphics/PiccoloVKDevice.h>
#include <Graphics/PiccoloVKFrameBuffer.h>

namespace Piccolo
{
    PiccoloVKRenderPass::PiccoloVKRenderPass(PiccoloVKDevice *device, const std::vector<Attachment> &attachments, const std::vector<RenderSubPass> &renderSubPasses)
        : mDevice(device), mAttachments(attachments), mRenderSubPasses(renderSubPasses)
    {
        // 默认subpass和附件
        if (mRenderSubPasses.empty())
        {
            mAttachments =
            {
                {
                    .format = device->GetSettings().surfaceFormat,
                    .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
                    .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
                    .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                    .usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT
                }
            };
            mRenderSubPasses =
            {
                { .colorAttachments = { 0 }, .sampleCount = VK_SAMPLE_COUNT_1_BIT }
            };
        }

        // 构建subpass
        std::vector<VkSubpassDescription> subPassDescriptions(mRenderSubPasses.size());
        std::vector<VkAttachmentReference> inputAttachmentRefs[mRenderSubPasses.size()];
        std::vector<VkAttachmentReference> colorAttachmentRefs[mRenderSubPasses.size()];
        std::vector<VkAttachmentReference> depthStencilAttachmentRefs[mRenderSubPasses.size()];
        VkAttachmentReference resolveAttachmentRefs[mRenderSubPasses.size()];

        for (int i = 0; i < mRenderSubPasses.size(); ++i)
        {
            RenderSubPass subPass = mRenderSubPasses[i];

            for (const auto& inputAttachment : subPass.inputAttachments)
            {
                inputAttachmentRefs[i].push_back({ inputAttachment, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL });
            }

            for (const auto& colorAttachment : subPass.colorAttachments)
            {
                colorAttachmentRefs[i].push_back({ colorAttachment, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL });
                mAttachments[colorAttachment].samples = subPass.sampleCount;
                if (subPass.sampleCount > VK_SAMPLE_COUNT_1_BIT)
                {
                    mAttachments[colorAttachment].finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
                }
            }

            for (const auto& depthStencilAttachment : subPass.depthStencilAttachments)
            {
                depthStencilAttachmentRefs[i].push_back({ depthStencilAttachment, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL });
                mAttachments[depthStencilAttachment].samples = subPass.sampleCount;
            }

            if (subPass.sampleCount > VK_SAMPLE_COUNT_1_BIT)
            {
                mAttachments.push_back({
                    .format = device->GetSettings().surfaceFormat,
                    .loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    .storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
                    .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                    .samples = VK_SAMPLE_COUNT_1_BIT,
                    .usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT
                });
                resolveAttachmentRefs[i] = { static_cast<uint32_t>(mAttachments.size() - 1), VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL };
            }

            subPassDescriptions[i].flags = 0;
            subPassDescriptions[i].pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
            subPassDescriptions[i].inputAttachmentCount = inputAttachmentRefs[i].size();
            subPassDescriptions[i].pInputAttachments = inputAttachmentRefs[i].data();
            subPassDescriptions[i].colorAttachmentCount = colorAttachmentRefs[i].size();
            subPassDescriptions[i].pColorAttachments = colorAttachmentRefs[i].data();
            subPassDescriptions[i].pResolveAttachments = subPass.sampleCount > VK_SAMPLE_COUNT_1_BIT ? &resolveAttachmentRefs[i] : nullptr;
            subPassDescriptions[i].pDepthStencilAttachment = depthStencilAttachmentRefs[i].data();
            subPassDescriptions[i].preserveAttachmentCount = 0;
            subPassDescriptions[i].pPreserveAttachments = nullptr;
        }


        std::vector<VkSubpassDependency> dependencies(mRenderSubPasses.size() - 1);
        if (mRenderSubPasses.size() > 1)
        {
            for (int i = 0; i < dependencies.size(); ++i)
            {
                dependencies[i].srcSubpass = i;
                dependencies[i].dstSubpass = i + 1;
                dependencies[i].srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
                dependencies[i].dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
                dependencies[i].srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
                dependencies[i].dstAccessMask = VK_ACCESS_INPUT_ATTACHMENT_READ_BIT;
                dependencies[i].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;
            }
        }

        // renderpass
        std::vector<VkAttachmentDescription> vkAttachments;
        vkAttachments.reserve(mAttachments.size());
        for (const auto& attachment : mAttachments)
        {
            vkAttachments.push_back({
                .flags = 0,
                .format = attachment.format,
                .samples = attachment.samples,
                .loadOp = attachment.loadOp,
                .storeOp = attachment.storeOp,
                .stencilLoadOp = attachment.stencilLoadOp,
                .stencilStoreOp = attachment.stencilStoreOp,
                .initialLayout = attachment.initialLayout,
                .finalLayout = attachment.finalLayout
            });
        }
        VkRenderPassCreateInfo renderPassCreateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .attachmentCount = static_cast<uint32_t>(vkAttachments.size()),
            .pAttachments = vkAttachments.data(),
            .subpassCount = static_cast<uint32_t>(mRenderSubPasses.size()),
            .pSubpasses = subPassDescriptions.data(),
            .dependencyCount = static_cast<uint32_t>(dependencies.size()),
            .pDependencies = dependencies.data()
        };

        CALL_VK(vkCreateRenderPass(device->GetLogicDevice(), &renderPassCreateInfo, nullptr, &mRenderPass));
        LOG_T("RenderPass {0} : {1}, attachment count: {2}, subpass count: {3}", __FUNCTION__, (void*)mRenderPass, mAttachments.size(), mRenderSubPasses.size());
    }

    PiccoloVKRenderPass::~PiccoloVKRenderPass()
    {
        VK_D(RenderPass, mDevice->GetLogicDevice(), mRenderPass);
    }

    void PiccoloVKRenderPass::BeginRenderPass(VkCommandBuffer commandBuffer, PiccoloVKFrameBuffer *frameBuffer, const std::vector<VkClearValue> &clearValues) const
    {
        VkRect2D renderArea =
        {
            .offset = { 0, 0 },
            .extent = { frameBuffer->GetWidth(), frameBuffer->GetHeight() }
        };
        VkRenderPassBeginInfo beginInfo =
        {
            .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
            .pNext = nullptr,
            .renderPass = mRenderPass,
            .framebuffer = frameBuffer->GetFrameBuffer(),
            .renderArea = renderArea,
            .clearValueCount = static_cast<uint32_t>(clearValues.size()),
            .pClearValues = clearValues.data()
        };
        vkCmdBeginRenderPass(commandBuffer, &beginInfo, VK_SUBPASS_CONTENTS_INLINE);
    }

    void PiccoloVKRenderPass::EndRenderPass(VkCommandBuffer commandBuffer) const
    {
        vkCmdEndRenderPass(commandBuffer);
    }
}
